import networkx as nx
import numpy as np
from collections import defaultdict
import logging
import torch
import torch.nn as nn
import torch.nn.functional as F
from sklearn.cluster import KMeans

# 尝试导入community模块
try:
    import community.community_louvain as community_louvain
except ImportError:
    try:
        from community import best_partition
        community_louvain = type('CommunityLouvain', (), {'best_partition': lambda g: best_partition(g)})
    except ImportError:
        print("警告: 无法导入community模块，将使用简单的替代算法")
        community_louvain = type('SimplePartition', (), {
            'best_partition': lambda g: {node: i for i, node in enumerate(g.nodes())}
        })

# 设置日志
logger = logging.getLogger("dhcds")
if not logger.handlers:
    handler = logging.StreamHandler()
    formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
    handler.setFormatter(formatter)
    logger.addHandler(handler)
    logger.setLevel(logging.INFO)

class DynamicHeterogeneousCommunityDetection:
    """动态异构图社区检测"""
    
    def __init__(self, stability_weight=0.5, embedding_dim=32, num_clusters=10):
        """初始化
        
        参数:
            stability_weight: 社区稳定性权重
            embedding_dim: 嵌入维度
            num_clusters: 默认社区数量
        """
        self.stability_weight = stability_weight
        self.embedding_dim = embedding_dim
        self.num_clusters = num_clusters
        
        self.previous_communities = None
        self.community_evolution = []
        self.node_type_weights = {}
        self.temporal_communities = []
        self.node_embeddings = {}
        self.community_embeddings = {}
    
    def _ensure_graph(self, data):
        """确保输入数据是NetworkX图对象，增强健壮性处理各种格式"""
        logger.info(f"_ensure_graph: 接收到数据类型 {type(data)}")
        
        # 已经是NetworkX图
        if isinstance(data, nx.Graph):
            return data
        
        # 创建新图
        G = nx.Graph()
        
        try:
            # 处理字典格式
            if isinstance(data, dict):
                # 标准格式 {'nodes': {...}, 'edges': [...]}
                if 'nodes' in data and isinstance(data['nodes'], dict):
                    for node_id, attrs in data['nodes'].items():
                        if isinstance(attrs, dict):
                            G.add_node(node_id, **attrs)
                        else:
                            G.add_node(node_id)
                    
                    if 'edges' in data and isinstance(data['edges'], list):
                        for edge in data['edges']:
                            if isinstance(edge, dict) and 'source' in edge and 'target' in edge:
                                attrs = {k: v for k, v in edge.items() if k not in ['source', 'target']}
                                G.add_edge(edge['source'], edge['target'], **attrs)
                            elif isinstance(edge, (list, tuple)) and len(edge) >= 2:
                                G.add_edge(edge[0], edge[1])
                
                # 另一种常见格式 {'vertices': [...], 'links': [...]}
                elif 'vertices' in data or 'links' in data:
                    # 处理节点
                    for node in data.get('vertices', []):
                        if isinstance(node, dict) and 'id' in node:
                            node_attrs = {k: v for k, v in node.items() if k != 'id'}
                            G.add_node(node['id'], **node_attrs)
                        elif isinstance(node, (str, int)):
                            G.add_node(node)
                    
                    # 处理边
                    for link in data.get('links', []):
                        if isinstance(link, dict) and 'source' in link and 'target' in link:
                            attrs = {k: v for k, v in link.items() if k not in ['source', 'target']}
                            G.add_edge(link['source'], link['target'], **attrs)
                
                # 邻接列表 {node: [neighbors...], ...}
                else:
                    for node, neighbors in data.items():
                        G.add_node(node)
                        if isinstance(neighbors, (list, tuple)):
                            for neighbor in neighbors:
                                if neighbor is not None:  # 确保邻居不是None
                                    G.add_edge(node, neighbor)
                        elif isinstance(neighbors, dict):
                            for neighbor, attr in neighbors.items():
                                if neighbor is not None:  # 确保邻居不是None
                                    if isinstance(attr, dict):
                                        G.add_edge(node, neighbor, **attr)
                                    else:
                                        G.add_edge(node, neighbor, weight=attr)
            
            # 处理边列表 [(source, target), ...]
            elif isinstance(data, (list, tuple)):
                for item in data:
                    if isinstance(item, (list, tuple)) and len(item) >= 2:
                        G.add_edge(item[0], item[1])
            
            # 处理其他格式，尝试添加基本节点
            else:
                # 尝试转换为字符串并添加为单个节点
                try:
                    G.add_node(str(data))
                except:
                    G.add_node(0)  # 添加一个默认节点
        
        except Exception as e:
            logger.error(f"图转换失败: {e}", exc_info=True)
            # 创建一个空图而不是抛出异常
            G = nx.Graph()
            G.add_node(0)  # 添加一个默认节点，避免空图
        
        logger.info(f"图转换完成，节点数: {G.number_of_nodes()}, 边数: {G.number_of_edges()}")
        return G
    
    def _simple_node_embeddings(self, graph):
        """简化版节点嵌入，不使用深度学习"""
        embeddings = {}
        dim = self.embedding_dim
        
        for node in graph.nodes():
            # 使用节点的简单特征作为嵌入
            degree = graph.degree(node)
            clustering = nx.clustering(graph, node)
            neighbors = list(graph.neighbors(node))
            num_neighbors = len(neighbors)
            
            # 创建基础特征
            features = [
                degree / (max(1, graph.number_of_nodes())),
                clustering,
                num_neighbors / (max(1, graph.number_of_nodes()))
            ]
            
            # 扩展到所需维度
            embedding = np.zeros(dim)
            for i in range(min(len(features), dim)):
                embedding[i] = features[i]
            
            # 添加随机噪声作为区分
            noise = np.random.randn(dim) * 0.1
            embedding += noise
            
            # 标准化
            norm = np.linalg.norm(embedding)
            if norm > 0:
                embedding = embedding / norm
            
            embeddings[node] = embedding
        
        self.node_embeddings = embeddings
        return embeddings
    
    def _extract_subgraph(self, graph, node_type):
        """提取特定类型节点的子图"""
        nodes = [n for n, attr in graph.nodes(data=True) if attr.get('type') == node_type]
        return graph.subgraph(nodes)
    
    def _find_type_communities(self, graph, node_type, prev_communities=None):
        """为特定类型的节点找到社区"""
        try:
            subgraph = self._extract_subgraph(graph, node_type)
            
            if len(subgraph) == 0:
                return {}
            
            # 使用Louvain算法
            try:
                partition = community_louvain.best_partition(subgraph)
            except Exception as e:
                logger.warning(f"Louvain社区检测失败: {e}，使用简单社区检测")
                # 使用简单的社区分配
                partition = {node: i % self.num_clusters for i, node in enumerate(subgraph.nodes())}
            
            # 如果有上一时刻的社区，考虑社区稳定性
            if prev_communities:
                # 计算重叠并调整
                self._apply_stability(partition, prev_communities)
            
            return partition
        except Exception as e:
            logger.exception(f"检测类型'{node_type}'的社区时出错: {e}")
            return {}
    
    def _apply_stability(self, partition, prev_communities):
        """应用社区稳定性约束"""
        # 计算新社区与旧社区的重叠
        overlaps = defaultdict(lambda: defaultdict(int))
        
        for node, new_comm in partition.items():
            if node in prev_communities:
                old_comm = prev_communities[node]
                overlaps[new_comm][old_comm] += 1
        
        # 根据重叠调整社区分配
        for node, new_comm in list(partition.items()):
            if node in prev_communities:
                old_comm = prev_communities[node]
                
                # 随机决定是否保持稳定性
                if np.random.random() < self.stability_weight:
                    partition[node] = old_comm
    
    def _merge_communities(self, type_communities):
        """简化版合并不同类型的社区"""
        all_communities = {}
        comm_id = 0
        
        # 为每种类型的社区分配全局ID
        for node_type, communities in type_communities.items():
            type_to_global = {}
            
            # 为每个局部社区分配全局ID
            for local_id in set(communities.values()):
                type_to_global[local_id] = comm_id
                comm_id += 1
            
            # 将节点映射到全局社区
            for node, local_id in communities.items():
                all_communities[node] = type_to_global[local_id]
        
        return all_communities
    
    def _track_evolution(self, current, previous):
        """跟踪社区演化(简化版)"""
        if not previous:
            return {'new': set(current.values())}
        
        # 按社区分组节点
        current_comms = defaultdict(set)
        previous_comms = defaultdict(set)
        
        for node, comm in current.items():
            current_comms[comm].add(node)
        
        for node, comm in previous.items():
            previous_comms[comm].add(node)
        
        # 跟踪事件
        evolution = {
            'continue': [],  # 继续存在的社区
            'new': set(),    # 新社区
            'dissolve': set() # 消失的社区
        }
        
        # 检测继续存在和新社区
        for curr_comm, curr_nodes in current_comms.items():
            best_match = None
            best_overlap = 0
            
            for prev_comm, prev_nodes in previous_comms.items():
                overlap = len(curr_nodes & prev_nodes)
                if overlap > best_overlap:
                    best_overlap = overlap
                    best_match = prev_comm
            
            # 判断社区延续还是新出现
            if best_match is not None and best_overlap / len(curr_nodes) > 0.3:
                evolution['continue'].append((best_match, curr_comm))
            else:
                evolution['new'].add(curr_comm)
        
        # 检测消失的社区
        for prev_comm in previous_comms:
            if not any(prev == prev_comm for prev, _ in evolution['continue']):
                evolution['dissolve'].add(prev_comm)
        
        return evolution
    
    def detect_communities(self, graph_data, previous_result=None):
        """检测异构图中的社区"""
        try:
            # 确保输入是图对象
            graph = self._ensure_graph(graph_data)
            
            # 空图检查
            if graph.number_of_nodes() == 0:
                logger.warning("图中没有节点，返回空结果")
                return {
                    'communities': {},
                    'type_communities': {},
                    'evolution': {'new': set()},
                    'community_embeddings': {}
                }
            
            # 学习节点嵌入(简化版)
            self._simple_node_embeddings(graph)
            
            # 获取节点类型
            node_types = set()
            for _, attr in graph.nodes(data=True):
                node_type = attr.get('type', 'unknown')
                node_types.add(node_type)
            
            # 为每种节点类型检测社区
            type_communities = {}
            for node_type in node_types:
                prev_comms = None
                if previous_result and 'type_communities' in previous_result:
                    prev_comms = previous_result['type_communities'].get(node_type)
                
                type_communities[node_type] = self._find_type_communities(
                    graph, node_type, prev_comms
                )
            
            # 合并社区
            merged_communities = self._merge_communities(type_communities)
            
            # 跟踪演化
            evolution = self._track_evolution(
                merged_communities,
                previous_result['communities'] if previous_result else None
            )
            
            # 计算社区嵌入(简化版)
            community_embeddings = self._calculate_community_embeddings(merged_communities)
            
            # 构建结果
            result = {
                'communities': merged_communities,
                'type_communities': type_communities,
                'evolution': evolution,
                'community_embeddings': community_embeddings
            }
            
            # 更新历史
            self.previous_communities = merged_communities
            self.temporal_communities.append(result)
            
            return result
        except Exception as e:
            logger.exception(f"社区检测失败: {e}")
            # 返回空结果而不是抛出异常
            return {
                'communities': {},
                'type_communities': {},
                'evolution': {'new': set()},
                'community_embeddings': {},
                'error': str(e)
            }
    
    def _calculate_community_embeddings(self, communities):
        """计算社区嵌入"""
        community_nodes = defaultdict(list)
        for node, comm in communities.items():
            community_nodes[comm].append(node)
        
        community_embeddings = {}
        for comm, nodes in community_nodes.items():
            # 收集节点嵌入
            node_embs = []
            for node in nodes:
                if node in self.node_embeddings:
                    node_embs.append(self.node_embeddings[node])
            
            # 计算平均嵌入
            if node_embs:
                comm_emb = np.mean(node_embs, axis=0)
                community_embeddings[comm] = comm_emb
        
        self.community_embeddings = community_embeddings
        return community_embeddings
    
    def detect_dynamic_communities(self, graph_sequence):
        """检测动态社区"""
        results = []
        previous_result = None
        
        # 检查输入格式
        if not isinstance(graph_sequence, (list, tuple)):
            if isinstance(graph_sequence, dict):
                # 单个图
                graph_sequence = [graph_sequence]
            else:
                logger.warning(f"无效的图序列类型: {type(graph_sequence)}")
                return []
        
        # 处理每个时间步
        for i, graph_data in enumerate(graph_sequence):
            logger.info(f"处理时间步 {i+1}/{len(graph_sequence)}")
            
            # 检测社区
            result = self.detect_communities(graph_data, previous_result)
            results.append(result)
            previous_result = result
            
            # 记录演化
            if i > 0:
                evolution = self._track_evolution(
                    result['communities'],
                    results[i-1]['communities']
                )
                result['evolution'] = evolution
        
        return results
    
    def get_community_members(self, community_id, time_step=-1):
        """获取社区成员"""
        if not self.temporal_communities:
            return []
        
        if time_step < 0:
            time_step = len(self.temporal_communities) + time_step
        
        if time_step < 0 or time_step >= len(self.temporal_communities):
            return []
        
        communities = self.temporal_communities[time_step]['communities']
        return [node for node, comm in communities.items() if comm == community_id]
    
    def calculate_community_similarity(self, comm1, time1, comm2, time2):
        """计算社区相似度"""
        members1 = self.get_community_members(comm1, time1)
        members2 = self.get_community_members(comm2, time2)
        
        # Jaccard相似度
        set1 = set(members1)
        set2 = set(members2)
        
        if not set1 or not set2:
            return 0.0
        
        intersection = len(set1 & set2)
        union = len(set1 | set2)
        
        return intersection / union

    @staticmethod
    def convert_to_networkx(data):
        """将各种数据格式转换为NetworkX图
        
        支持以下格式:
        1. NetworkX图对象
        2. 字典格式的图表示 {'nodes': {...}, 'edges': [...]}
        3. 邻接列表 {node: [neighbors...], ...}
        4. 边列表 [(source, target), ...]
        5. 节点和边列表 {'vertices': [...], 'links': [...]}
        """
        logger.info(f"convert_to_networkx: 接收到数据类型 {type(data)}")
        
        # 已经是NetworkX图
        if isinstance(data, nx.Graph):
            return data
        
        # 创建一个新图
        G = nx.Graph()
        
        # 字典格式
        if isinstance(data, dict):
            # 检查是否为标准格式
            if 'nodes' in data or 'edges' in data:
                # 标准格式 {'nodes': {...}, 'edges': [...]}
                for node_id, attrs in data.get('nodes', {}).items():
                    if isinstance(attrs, dict):
                        G.add_node(node_id, **attrs)
                    else:
                        G.add_node(node_id)
                        
                for edge in data.get('edges', []):
                    if isinstance(edge, dict) and 'source' in edge and 'target' in edge:
                        attrs = {k: v for k, v in edge.items() if k not in ['source', 'target']}
                        G.add_edge(edge['source'], edge['target'], **attrs)
                    elif isinstance(edge, (list, tuple)) and len(edge) >= 2:
                        if len(edge) > 2 and isinstance(edge[2], dict):
                            G.add_edge(edge[0], edge[1], **edge[2])
                        else:
                            G.add_edge(edge[0], edge[1])
                            
            # 检查是否为另一种常见格式
            elif 'vertices' in data or 'links' in data:
                # 另一种常见格式 {'vertices': [...], 'links': [...]}
                for node in data.get('vertices', []):
                    if isinstance(node, dict) and 'id' in node:
                        node_id = node.pop('id')
                        G.add_node(node_id, **node)
                    elif isinstance(node, (str, int)):
                        G.add_node(node)
                        
                for link in data.get('links', []):
                    if isinstance(link, dict) and 'source' in link and 'target' in link:
                        attrs = {k: v for k, v in link.items() if k not in ['source', 'target']}
                        G.add_edge(link['source'], link['target'], **attrs)
                        
            # 检查是否为邻接列表
            else:
                # 假设是邻接列表 {node: [neighbors...], ...}
                for node, neighbors in data.items():
                    G.add_node(node)
                    if isinstance(neighbors, (list, tuple, set)):
                        for neighbor in neighbors:
                            G.add_edge(node, neighbor)
                    elif isinstance(neighbors, dict):
                        for neighbor, edge_attr in neighbors.items():
                            if isinstance(edge_attr, dict):
                                G.add_edge(node, neighbor, **edge_attr)
                            else:
                                G.add_edge(node, neighbor, weight=edge_attr)
                                
        # 边列表
        elif isinstance(data, (list, tuple)):
            for item in data:
                if isinstance(item, (list, tuple)) and len(item) >= 2:
                    if len(item) > 2 and isinstance(item[2], dict):
                        G.add_edge(item[0], item[1], **item[2])
                    else:
                        G.add_edge(item[0], item[1])
        
        # 如果图是空的，记录警告
        if G.number_of_nodes() == 0:
            logger.warning("转换结果是一个空图")
        
        return G 
